// Package ais provides AIStore's proxy and target nodes.
/*
 * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 */
package ais

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"net/http"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/NVIDIA/aistore/cmn/cos"
	"github.com/NVIDIA/aistore/cmn/debug"
	"github.com/NVIDIA/aistore/cmn/mono"
	"github.com/NVIDIA/aistore/cmn/nlog"
	"github.com/NVIDIA/aistore/memsys"
)

// This source contains cluster-wide HMAC signing for HTTP redirect URLs.
//
// Cluster Shared Key (aka "Cluster Secret Key" or `clusterKey`):
//   A cryptographically random secret generated by the primary proxy
//   and distributed to all cluster nodes via metasync. The key is versioned and tied to the
//   current Smap (cluster map) version for consistency during topology changes.
// Other stateful entities include:
//   - cskOwner: holds {secret, ver, created} + nonce
//   - signer:   populates its own state and computes over HMAC payload
// Configuration:
//   Controlled via cluster config: auth.cluster_key.enabled
//   See cmn.ClusterKeyConf for time windows and related settings.
// HMAC:
//   When enabled, proxies sign redirect URLs using HMAC-SHA256. The signature currently covers:
//   HTTP method, URL path, proxy ID, Smap version, content-length, and a monotonic nonce.
//   Incoming requests are validated once at the datapath entry point (h.parseReq => sign.verify()):
//     - dpq fast-path embeds cskgrp directly;
//     - url.Values slow-path uses cskFromQ().
//   signer.verify() reconstructs the HMAC payload (using pid, smapVer, nonce, and URL fields)
//   and compares it against the provided signature. Any mismatch results in a 401.

// TODO: sign.verify()
// - check that csk is always initialized across (enable/disable; lifecycle events)
//   (see related sign.warn())
// - anti-replay logic (sliding time window, per-sender nonce tracking, etc.)
// - validate pid: a) != "" (weak) or b) in Smap (strong)
// - optionally, extend HMAC payload to cover assorted query parameters

const (
	cskTag = "csk"

	cskKeyLen      = 16               // len(secret)
	cskSigLen      = 43               // = base64.RawURLEncoding.EncodedLen(sha256.Size)
	cskSepa        = 0                // to separate strings in HMAC payload
	cskBase        = 36               // base for all signed int64/uint64 fields
	cskURLOverhead = 160              // pid+utm+vpams+x+u qparams (~145 worst-case)
	cskUptime      = 10 * time.Second // cluster uptime after which we start warning if CSK version remains zero
)

type (
	signer struct {
		r       *http.Request
		h       *htrun
		sb      *cos.SB
		sig     []byte
		smapVer int64
		nonce   uint64
	}
	cskOwner struct {
		k     atomic.Value // => clusterKey
		nonce atomic.Uint64
	}
	clusterKey struct {
		secret  []byte
		ver     int64
		created int64
	}
	// hash and buffer
	handb struct {
		h   hash.Hash
		ver int64
		buf [sha256.Size]byte
	}
)

// interface guard
var (
	_ cos.Packer   = (*clusterKey)(nil)
	_ cos.Unpacker = (*clusterKey)(nil)
	_ revs         = (*clusterKey)(nil)
)

//////////////
// cskOwner //
//////////////

func (csk *cskOwner) init() {
	i := cos.CryptoRandI()
	csk.nonce.Store(uint64(i))
	csk.reset()
}

func (csk *cskOwner) load() (k *clusterKey) { return csk.k.Load().(*clusterKey) }
func (csk *cskOwner) store(k *clusterKey)   { csk.k.Store(k) }
func (csk *cskOwner) reset()                { csk.k.Store(&clusterKey{}) }

// primary only
// version is monotonically increasing and is loosely tied to smap version:
// the latter is strictly guarded by primary and  can therefore, be relied
// upon in re: false-positive downgrades
func (csk *cskOwner) gen(smapVer int64) (nk *clusterKey) {
	ok := csk.load()
	nk = &clusterKey{
		secret:  cos.CryptoRandB(cskKeyLen),
		ver:     max(smapVer, ok.ver+1),
		created: mono.NanoTime(),
	}
	csk.store(nk)
	return
}

////////////
// signer //
////////////

func (sign *signer) bufsize(pid string) int {
	r := sign.r
	return len(r.Method) + 1 + len(r.URL.Path) + 1 + len(pid) + 1 + 3*cos.SizeofI64 + cskSigLen
}

func (sign *signer) compute(pid string) {
	var (
		r    = sign.r
		sb   = sign.sb
		size = sign.bufsize(pid)
	)
	sb.Reset(size, false /*allow shrink*/)
	debug.Assert(sb.Cap() >= size, sb.Cap(), " vs ", size)

	// (method, url path, pid)
	sb.WriteString(r.Method)
	sb.WriteUint8(cskSepa)
	sb.WriteString(r.URL.Path)
	sb.WriteUint8(cskSepa)
	sb.WriteString(pid)
	sb.WriteUint8(cskSepa)

	// (smap, content-length, nonce)
	var b8 [8]byte
	binary.BigEndian.PutUint64(b8[:], uint64(sign.smapVer))
	sb.WriteBytes(b8[:])
	binary.BigEndian.PutUint64(b8[:], uint64(max(r.ContentLength, 0))) // (not max uint64)
	sb.WriteBytes(b8[:])
	binary.BigEndian.PutUint64(b8[:], sign.nonce)
	sb.WriteBytes(b8[:])

	// load key
	k := sign.h.owner.csk.load()
	if k.ver == 0 {
		sign.warn()
		return
	}
	debug.Assert(len(k.secret) > 0, k.String())

	// hash
	hb := handAlloc()
	if hb.h == nil || hb.ver != k.ver {
		hb.h = hmac.New(sha256.New, k.secret)
		hb.ver = k.ver
	} else {
		hb.h.Reset()
	}
	buf, l := sb.Bytes(), sb.Len()
	hb.h.Write(buf[:l])
	hb.h.Sum(hb.buf[:0])

	// encode and append => sb
	sign.sig = sb.ReserveAppend(cskSigLen)
	base64.RawURLEncoding.Encode(sign.sig, hb.buf[:])

	handFree(hb)
}

func (sign *signer) warn() {
	const skip = "skip HMAC sign (cluster-key version is zero)"
	uptime := sign.h.keepalive.cluUptime(mono.NanoTime())
	switch {
	case uptime < cskUptime:
		// expected during early startup, stay quiet
	case uptime < cskUptime<<1:
		nlog.Warningln(sign.h.String(), skip)
	case uptime < min(time.Minute, cskUptime<<2):
		nlog.Errorln(sign.h.String(), skip)
	}
}

func (sign *signer) buildURL(nodeURL string, now int64) string {
	var (
		h  = sign.h
		r  = sign.r
		sb = sign.sb
	)
	// 1) build the query string while `sign.sig` still points into sb.buf
	q := qAlloc()
	raw := h.qencode(q, now, sign)
	qFree(q)

	// 2) reuse sb for final URL
	size := len(nodeURL) + len(r.URL.Path) + len(r.URL.RawQuery) + len(raw) + 2
	debug.Assert(sb.Cap() >= size, sb.Cap(), " vs ", size)
	sb.Reset(size, false)

	sb.WriteString(nodeURL)
	sb.WriteString(r.URL.Path)
	sb.WriteUint8('?')

	if r.URL.RawQuery != "" {
		sb.WriteString(r.URL.RawQuery)
		sb.WriteUint8('&')
	}

	sb.WriteString(raw)
	return sb.CloneString()
}

// receive-side HMAC validation
// on failure return http.StatusUnauthorized or 0 (that is, 400), the latter for invalid request
func (sign *signer) verify(pid string, cskgrp *cskgrp) (int, error) {
	if len(cskgrp.hmacSig) != cskSigLen {
		return 0, fmt.Errorf("invalid signature length: %d", len(cskgrp.hmacSig))
	}
	sign.nonce = cskgrp.nonce
	sign.smapVer = cskgrp.smapVer

	size := sign.bufsize(pid)
	sign.sb = sbAlloc()
	sign.sb.Reset(size, true /*allow shrink*/)

	sign.compute(pid)
	sbFree(sign.sb)

	if !bytes.Equal(sign.sig, cos.UnsafeB(cskgrp.hmacSig)) {
		return http.StatusUnauthorized, errors.New("HMAC signature mismatch")
	}
	return 0, nil
}

////////////////
// clusterKey //
////////////////

// as byte-packer
func (k *clusterKey) PackedSize() int {
	return cos.SizeofI64 + cos.SizeofI64 + cos.PackedBytesLen(k.secret)
}

func (k *clusterKey) Pack(packer *cos.BytePack) {
	packer.WriteInt64(k.ver)
	packer.WriteInt64(k.created)
	packer.WriteBytes(k.secret)
}

func (k *clusterKey) Unpack(unpacker *cos.ByteUnpack) (err error) {
	if k.ver, err = unpacker.ReadInt64(); err != nil {
		return err
	}
	if k.created, err = unpacker.ReadInt64(); err != nil {
		return err
	}
	k.secret, err = unpacker.ReadBytes()
	return err
}

// as metasync `revs`
func (*clusterKey) tag() string       { return revsCSKTag }
func (k *clusterKey) version() int64  { return k.ver }
func (*clusterKey) uuid() string      { return "" }
func (k *clusterKey) jit(*proxy) revs { return k }
func (*clusterKey) sgl() *memsys.SGL  { return nil }

func (k *clusterKey) marshal() []byte {
	size := k.PackedSize()
	packer := cos.NewPacker(make([]byte, size), size)
	packer.WriteAny(k)
	return packer.Bytes()
}

func (k *clusterKey) String() string {
	return cskTag + " v" + strconv.FormatInt(k.ver, 10)
}

//
// mem-pools
//

var (
	handPool = sync.Pool{New: func() any { return new(handb) }}  // reusable (versioned) HMAC hashes
	sbPool   = sync.Pool{New: func() any { return new(cos.SB) }} // reusable buffers
)

func handAlloc() *handb     { return handPool.Get().(*handb) }
func handFree(handb *handb) { handb.h.Reset(); handPool.Put(handb) }

func sbAlloc() *cos.SB  { return sbPool.Get().(*cos.SB) }
func sbFree(sb *cos.SB) { sbPool.Put(sb) }
